import torch
import torch.nn as nn
import torch.nn.functional as F
from ghost_module import GhostModule
from hornet_block import HornetLiteBlock

class Bottleneck(nn.Module):
    def __init__(self, in_planes, planes, stride=1, use_hornet=False):
        super(Bottleneck, self).__init__()
        self.use_hornet = use_hornet
        self.ghost1 = GhostModule(in_planes, planes, kernel_size=1)
        self.ghost2 = GhostModule(planes, planes, kernel_size=3, stride=stride)
        if use_hornet:
            self.hornet = HornetLiteBlock(planes)
        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes)
            )

    def forward(self, x):
        out = self.ghost1(x)
        out = self.ghost2(out)
        if self.use_hornet:
            out = self.hornet(out)
        out += self.shortcut(x)
        return F.relu(out)
